import paddle
import numpy
import os
from data_loader import Mydata
from paddle.vision.transforms import Normalize,Compose,Resize
from paddle.io import SequenceSampler, RandomSampler, BatchSampler, DistributedBatchSampler
from paddle.static import InputSpec
from paddle.jit import to_static
from paddle.vision.transforms import ToTensor

# base_dir='C:/Users/86167/Desktop/program/garbage'
# train_data_dir=os.path.join(base_dir,'train_data')
# train_label_dir=os.path.join(base_dir,'train_label')
# test_data_dir=os.path.join(base_dir,'test_data')
# test_label_dir=os.path.join(base_dir,'test_label')
if __name__=="__main__":
    transform = Compose([Resize((224,224)),Normalize(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], data_format='CHW')])
    # # transform=Normalize(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], data_format='CHW')
    # train_custom_dataset = Mydata(train_data_dir, train_label_dir, transform)


    # sampler = RandomSampler(train_custom_dataset)
    # train_custom_dataset = BatchSampler(sampler=sampler, batch_size=10)
    # # print(train_custom_dataset)
    # for data in train_custom_dataset:
    #     print(data)
    #
    # test_custom_dataset = Mydata(test_data_dir, test_label_dir, transform)
    # print('train_custom_dataset images: ', len(train_custom_dataset), 'test_custom_dataset images: ',
    #       len(test_custom_dataset))
    #
    # print("飞桨的内置模型：", paddle.vision.models.__all__)
    # model = paddle.vision.models.vgg16(batch_norm=True)
    # model = paddle.Model(mnist)
    # 为模型训练做准备，设置优化器及其学习率，并将网络的参数传入优化器，设置损失函数和精度计算方式
    import paddle
    from paddle.vision.transforms import ToTensor
    from paddle.vision.models import vgg16
    print("模型建立！\t")
    # build model
    model = vgg16()

    # build vgg16 model with batch_norm
    model = vgg16(batch_norm=True)

    # 使用高层API——paddle.Model对模型进行封装
    model = paddle.Model(model)
    # print("模型准备！\t")
    # # 为模型训练做准备，设置优化器，损失函数和精度计算方式
    # model.prepare(optimizer=paddle.optimizer.Adam(parameters=model.parameters()),
    #               loss=paddle.nn.CrossEntropyLoss(),
    #               metrics=paddle.metric.Accuracy())
    # print("模型训练启动！\t")
    # # 启动模型训练，指定训练数据集，设置训练轮次，设置每次数据集计算的批次大小，设置日志格式
    # model.fit(train_custom_dataset,
    #           test_custom_dataset,
    #           epochs=10,
    #           batch_size=10,
    #           save_dir="vgg16/",
    #           save_freq=10,
    #           verbose=1)
    # paddle.save(model.state_dict(), 'lenet.pdparams')
    # paddle.save(optim.state_dict(), "lenet.pdopt")
    # 转化模型为静态图，方便推理模式
    # ——————————————————————————————————————————————————————————————#
    # ——————————————————————————————————————————————————————————————#
    # ——————————————————————————————————————————————————————————————#
    # ——————————————————————————————————————————————————————————————#
    # ——————————————————————————————————————————————————————————————#
    # ——————————————————————————————————————————————————————————————#
    # load training format model
    # model_state_dict = paddle.load('final.pdparams')
    # # opt_state_dict = paddle.load('lenet.pdopt')
    # model.set_state_dict(model_state_dict)
    # # optim.set_state_dict(opt_state_dict)
    #
    # # save inferencing format model
    # net = to_static(model,
    #                 input_spec=[InputSpec(shape=[None, 1, 28, 28], name='image')])
    # paddle.jit.save(net, 'vgg16')


    from tkinter import *
    import cv2
    from PIL import Image, ImageTk
    import os
    import numpy as np
    import matplotlib.pylab as plt

    # image_path = "C:/Users/86167/Desktop/program/垃圾分类/Garbage_classification/garbage/garbage/cardboard2.jpg"
    # image=Image.open(image_path)
    # plt.imshow(image)
    # plt.show()
    # print(image)
    # image=image.astype('float32')
    # print(image)

    mean_values = paddle.to_tensor([127.5, 127.5, 127.5])  # 这里的数值需要根据你的实际情况来设定
    std_values = paddle.to_tensor([127.5, 127.5, 127.5])  # 同样的，这里的数值也需要根据实际情况来设定
    transform = Compose([Resize((224, 224)), Normalize(mean_values, std_values, data_format='CHW')])
    print(mean_values.shape,std_values.shape)
    # model.load('final.pdparams')
    model.load('final.pdparams')
    image_path = "C:/Users/86167/Desktop/cardboard2.jpg"
    image = cv2.imread(image_path,cv2.IMREAD_GRAYSCALE)
    print(image.shape)
    if image is None:
        print(f"Can't open image at {image_path}")
    image = image.astype('float32')
    print(image.shape)
    if transform is not None:
        img = transform(image)

    model.predict(image)